#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
QTorch量化交易框架 - 参数优化演示
展示如何使用参数优化框架自动寻找最佳策略参数
"""

import logging
import argparse
import os
import time
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from datetime import datetime

from qtorch.core.qoptimizer import QOptimizer
from qtorch.strategy.simple_sma_strategy import SimpleSMAStrategy
from qtorch.strategy.simple_macd_strategy import SimpleMACDStrategy
from qtorch.strategy.enhanced_rsi_strategy import EnhancedRSIStrategy
from qtorch.strategy.multi_indicator_strategy import MultiIndicatorStrategy

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler()]
)

def optimize_simple_sma(symbol='000001', start_date='2020-01-01', end_date='2022-12-31', 
                     n_jobs=1, output_dir='./output/optimize'):
    """双均线策略参数优化示例"""
    logging.info("开始双均线策略参数优化...")
    
    # 创建优化器
    optimizer = QOptimizer(
        strategy_class=SimpleSMAStrategy,
        symbols=[symbol],
        start_date=start_date,
        end_date=end_date,
        output_dir=output_dir
    )
    
    # 定义参数网格
    param_grid = {
        'fast_period': list(range(5, 51, 5)),
        'slow_period': list(range(20, 201, 20))
    }
    
    # 运行网格搜索
    best_params = optimizer.grid_search(
        param_grid=param_grid,
        metric='夏普比率',
        higher_is_better=True,
        n_jobs=n_jobs,
        verbose=True
    )
    
    logging.info(f"双均线策略最佳参数: {best_params}")
    
    return best_params

def optimize_macd(symbol='000001', start_date='2020-01-01', end_date='2022-12-31', 
               n_jobs=1, output_dir='./output/optimize'):
    """MACD策略参数优化示例"""
    logging.info("开始MACD策略参数优化...")
    
    # 创建优化器
    optimizer = QOptimizer(
        strategy_class=SimpleMACDStrategy,
        symbols=[symbol],
        start_date=start_date,
        end_date=end_date,
        output_dir=output_dir
    )
    
    # 定义参数网格
    param_grid = {
        'fast_period': [8, 10, 12, 14, 16],
        'slow_period': [20, 24, 26, 28, 32],
        'signal_period': [7, 8, 9, 10, 11]
    }
    
    # 运行网格搜索
    best_params = optimizer.grid_search(
        param_grid=param_grid,
        metric='夏普比率',
        higher_is_better=True,
        n_jobs=n_jobs,
        verbose=True
    )
    
    logging.info(f"MACD策略最佳参数: {best_params}")
    
    return best_params

def optimize_enhanced_rsi(symbol='000001', start_date='2020-01-01', end_date='2022-12-31', 
                        n_jobs=1, output_dir='./output/optimize'):
    """增强RSI策略参数优化示例"""
    logging.info("开始增强RSI策略参数优化...")
    
    # 创建优化器
    optimizer = QOptimizer(
        strategy_class=EnhancedRSIStrategy,
        symbols=[symbol],
        start_date=start_date,
        end_date=end_date,
        output_dir=output_dir
    )
    
    # 使用随机搜索
    param_distributions = {
        'period': lambda: np.random.randint(10, 21),
        'overbought': lambda: np.random.randint(65, 85),
        'oversold': lambda: np.random.randint(15, 35),
        'use_ma_filter': lambda: np.random.choice([True, False]),
        'ma_period': lambda: np.random.randint(20, 101),
        'use_trend_filter': lambda: np.random.choice([True, False]),
        'use_fixed_exit': lambda: np.random.choice([True, False])
    }
    
    # 运行随机搜索
    best_params = optimizer.random_search(
        param_distributions=param_distributions,
        n_iter=30,
        metric='夏普比率',
        higher_is_better=True,
        n_jobs=n_jobs,
        verbose=True
    )
    
    logging.info(f"增强RSI策略最佳参数: {best_params}")
    
    return best_params

def optimize_multi_indicator(symbol='000001', start_date='2020-01-01', end_date='2022-12-31', 
                          n_jobs=1, output_dir='./output/optimize'):
    """多指标组合策略参数优化示例"""
    logging.info("开始多指标组合策略参数优化...")
    
    # 创建优化器
    optimizer = QOptimizer(
        strategy_class=MultiIndicatorStrategy,
        symbols=[symbol],
        start_date=start_date,
        end_date=end_date,
        output_dir=output_dir
    )
    
    # 使用随机搜索
    param_distributions = {
        'weight_ma': lambda: np.random.uniform(0.5, 2.0),
        'weight_rsi': lambda: np.random.uniform(0.5, 2.0),
        'weight_macd': lambda: np.random.uniform(0.5, 2.0),
        'weight_bb': lambda: np.random.uniform(0.5, 2.0),
        'threshold': lambda: np.random.uniform(0.3, 0.7)
    }
    
    # 运行随机搜索
    best_params = optimizer.random_search(
        param_distributions=param_distributions,
        n_iter=30,
        metric='夏普比率',
        higher_is_better=True,
        n_jobs=n_jobs,
        verbose=True
    )
    
    logging.info(f"多指标组合策略最佳参数: {best_params}")
    
    return best_params

def run_multi_symbol_optimization(strategy_class, symbols, param_grid, start_date, end_date,
                               metric='夏普比率', n_jobs=1, output_dir='./output/optimize_multi'):
    """多股票参数优化示例"""
    logging.info(f"开始多股票参数优化: {strategy_class.__name__}, {len(symbols)}个股票...")
    
    # 创建优化器
    optimizer = QOptimizer(
        strategy_class=strategy_class,
        symbols=symbols,  # 传入多个股票
        start_date=start_date,
        end_date=end_date,
        output_dir=output_dir
    )
    
    # 运行网格搜索
    best_params = optimizer.grid_search(
        param_grid=param_grid,
        metric=metric,
        higher_is_better=True,
        n_jobs=n_jobs,
        verbose=True
    )
    
    logging.info(f"多股票优化最佳参数: {best_params}")
    
    return best_params

def compare_optimized_strategies(start_date='2020-01-01', end_date='2022-12-31',
                              symbol='000001', output_dir='./output/compare'):
    """比较不同优化后的策略表现"""
    from qtorch.core.qengine import QEngine
    import matplotlib.pyplot as plt
    import os
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 优化各个策略，获取最佳参数
    # 这里使用简化的参数网格以加快示例运行速度
    sma_params = {'fast_period': 10, 'slow_period': 40}  # 简化示例
    macd_params = {'fast_period': 12, 'slow_period': 26, 'signal_period': 9}  # 简化示例
    rsi_params = {'period': 14, 'overbought': 70, 'oversold': 30}  # 简化示例
    
    # 创建策略列表
    strategies = [
        (SimpleSMAStrategy, sma_params, "双均线"),
        (SimpleMACDStrategy, macd_params, "MACD"),
        (EnhancedRSIStrategy, rsi_params, "RSI")
    ]
    
    # 运行各个策略并收集结果
    results = {}
    metrics = {}
    
    for strategy_class, params, name in strategies:
        # 创建引擎
        engine = QEngine()
        engine.add_data(symbol, start_date, end_date)
        engine.add_strategy(strategy_class, **params)
        
        # 运行回测
        logging.info(f"运行策略: {name}")
        result = engine.run()
        
        # 保存结果
        results[name] = result.data
        metrics[name] = result.metrics
        
        # 保存单个策略报告
        result.plot(output_dir=f"{output_dir}/{name}")
    
    # 比较各个策略的累积收益
    plt.figure(figsize=(12, 6))
    
    for name, data in results.items():
        if 'cum_returns' in data.columns:
            data['cum_returns'].plot(label=name)
    
    plt.title("策略优化对比 - 累积收益")
    plt.xlabel("日期")
    plt.ylabel("累积收益")
    plt.legend()
    plt.grid(True)
    plt.savefig(f"{output_dir}/strategies_comparison.png", dpi=150)
    
    # 创建关键指标对比表格
    metrics_comparison = {}
    for name, metric in metrics.items():
        metrics_comparison[name] = {
            "总收益率": metric.get("总收益率", 0),
            "年化收益率": metric.get("年化收益率", 0),
            "最大回撤": metric.get("最大回撤", 0),
            "夏普比率": metric.get("夏普比率", 0),
            "索提诺比率": metric.get("索提诺比率", 0),
            "卡玛比率": metric.get("卡玛比率", 0),
            "胜率": metric.get("胜率", 0)
        }
    
    # 转换为DataFrame
    df_comparison = pd.DataFrame(metrics_comparison)
    
    # 保存对比结果
    comparison_path = f"{output_dir}/metrics_comparison.csv"
    df_comparison.to_csv(comparison_path)
    logging.info(f"策略对比结果已保存至: {comparison_path}")
    
    # 打印对比结果
    print("\n===== 策略优化对比 =====")
    print(df_comparison)
    
    return df_comparison

def demonstrate_risk_management(symbol='000001', start_date='2020-01-01', end_date='2022-12-31',
                              output_dir='./output/risk_management'):
    """演示风险管理参数调整功能"""
    from qtorch.core.qengine import QEngine
    from qtorch.core.riskmanager import RiskManager
    import matplotlib.pyplot as plt
    import os
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    logging.info("开始风险管理参数调整演示...")
    
    # 创建引擎和策略
    engine = QEngine()
    engine.add_data(symbol, start_date, end_date)
    engine.add_strategy(SimpleSMAStrategy, fast_period=10, slow_period=40)
    
    # 获取风险管理器实例
    risk_manager = RiskManager(
        max_position=0.1,
        max_loss=0.05,
        max_drawdown=0.2,
        volatility_limit=0.03
    )
    
    # 演示1: 使用adjust_parameters方法（使用max_dd替代max_drawdown）
    logging.info("演示1: 使用adjust_parameters方法调整风险参数")
    risk_manager.adjust_parameters(
        max_position=0.15,
        max_loss=0.06,
        max_dd=0.18,  # 注意这里使用max_dd而不是max_drawdown
        volatility_limit=0.035
    )
    
    # 打印调整后的参数
    logging.info(f"调整后的最大持仓: {risk_manager.max_position}")
    logging.info(f"调整后的最大亏损: {risk_manager.max_loss}")
    logging.info(f"调整后的最大回撤: {risk_manager.max_drawdown}")
    logging.info(f"调整后的波动率限制: {risk_manager.volatility_limit}")
    
    # 演示2: 使用apply_dynamic_parameters方法根据市场状态动态调整参数
    logging.info("\n演示2: 使用apply_dynamic_parameters方法动态调整风险参数")
    
    # 模拟不同市场状态下的参数调整
    market_states = ["normal", "volatile", "crisis"]
    risk_levels = ["low", "medium", "high"]
    
    # 创建结果表格
    results = []
    
    for market_state in market_states:
        for risk_level in risk_levels:
            # 重置风险管理器
            risk_manager = RiskManager(
                max_position=0.1,
                max_loss=0.05,
                max_drawdown=0.2,
                volatility_limit=0.03
            )
            
            # 应用动态参数调整
            risk_manager.apply_dynamic_parameters(
                market_state=market_state,
                volatility_factor=1.2 if market_state == "volatile" else (1.5 if market_state == "crisis" else 1.0),
                risk_level=risk_level
            )
            
            # 记录结果
            results.append({
                "市场状态": market_state,
                "风险等级": risk_level,
                "最大持仓": risk_manager.max_position,
                "最大亏损": risk_manager.max_loss,
                "最大回撤": risk_manager.max_drawdown,
                "波动率限制": risk_manager.volatility_limit
            })
    
    # 转换为DataFrame并保存
    df_results = pd.DataFrame(results)
    results_path = f"{output_dir}/dynamic_risk_params.csv"
    df_results.to_csv(results_path, index=False)
    logging.info(f"动态风险参数调整结果已保存至: {results_path}")
    
    # 可视化结果
    plt.figure(figsize=(12, 8))
    
    # 绘制不同市场状态下的最大回撤参数
    pivot_dd = df_results.pivot(index="风险等级", columns="市场状态", values="最大回撤")
    ax1 = plt.subplot(2, 2, 1)
    pivot_dd.plot(kind='bar', ax=ax1)
    ax1.set_title("不同市场状态下的最大回撤参数")
    ax1.set_ylabel("最大回撤")
    
    # 绘制不同市场状态下的最大持仓参数
    pivot_pos = df_results.pivot(index="风险等级", columns="市场状态", values="最大持仓")
    ax2 = plt.subplot(2, 2, 2)
    pivot_pos.plot(kind='bar', ax=ax2)
    ax2.set_title("不同市场状态下的最大持仓参数")
    ax2.set_ylabel("最大持仓")
    
    # 绘制不同市场状态下的最大亏损参数
    pivot_loss = df_results.pivot(index="风险等级", columns="市场状态", values="最大亏损")
    ax3 = plt.subplot(2, 2, 3)
    pivot_loss.plot(kind='bar', ax=ax3)
    ax3.set_title("不同市场状态下的最大亏损参数")
    ax3.set_ylabel("最大亏损")
    
    # 绘制不同市场状态下的波动率限制参数
    pivot_vol = df_results.pivot(index="风险等级", columns="市场状态", values="波动率限制")
    ax4 = plt.subplot(2, 2, 4)
    pivot_vol.plot(kind='bar', ax=ax4)
    ax4.set_title("不同市场状态下的波动率限制参数")
    ax4.set_ylabel("波动率限制")
    
    plt.tight_layout()
    plt.savefig(f"{output_dir}/dynamic_risk_params.png", dpi=150)
    
    logging.info("风险管理参数调整演示完成")
    return df_results

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='QTorch量化交易框架 - 参数优化演示')
    parser.add_argument('--mode', type=str, default='sma',
                        choices=['sma', 'macd', 'rsi', 'multi', 'compare', 'risk', 'all'],
                        help='优化模式')
    parser.add_argument('--symbol', type=str, default='000001',
                        help='股票代码')
    parser.add_argument('--start', type=str, default='2020-01-01',
                        help='回测开始日期')
    parser.add_argument('--end', type=str, default='2022-12-31',
                        help='回测结束日期')
    parser.add_argument('--jobs', type=int, default=1,
                        help='并行任务数')
    parser.add_argument('--output', type=str, default='./output/optimize',
                        help='输出目录')
                        
    args = parser.parse_args()
    
    # 创建输出目录
    os.makedirs(args.output, exist_ok=True)
    
    try:
        if args.mode == 'sma' or args.mode == 'all':
            optimize_simple_sma(
                symbol=args.symbol,
                start_date=args.start,
                end_date=args.end,
                n_jobs=args.jobs,
                output_dir=f"{args.output}/sma"
            )
            
        if args.mode == 'macd' or args.mode == 'all':
            optimize_macd(
                symbol=args.symbol,
                start_date=args.start,
                end_date=args.end,
                n_jobs=args.jobs,
                output_dir=f"{args.output}/macd"
            )
            
        if args.mode == 'rsi' or args.mode == 'all':
            optimize_enhanced_rsi(
                symbol=args.symbol,
                start_date=args.start,
                end_date=args.end,
                n_jobs=args.jobs,
                output_dir=f"{args.output}/rsi"
            )
            
        if args.mode == 'multi' or args.mode == 'all':
            optimize_multi_indicator(
                symbol=args.symbol,
                start_date=args.start,
                end_date=args.end,
                n_jobs=args.jobs,
                output_dir=f"{args.output}/multi"
            )
            
        if args.mode == 'compare' or args.mode == 'all':
            compare_optimized_strategies(
                start_date=args.start,
                end_date=args.end,
                symbol=args.symbol,
                output_dir=f"{args.output}/compare"
            )
            
        if args.mode == 'risk' or args.mode == 'all':
            demonstrate_risk_management(
                symbol=args.symbol,
                start_date=args.start,
                end_date=args.end,
                output_dir=f"{args.output}/risk_management"
            )
            
        logging.info(f"所有优化任务已完成！结果保存在: {args.output}")
            
    except Exception as e:
        logging.error(f"优化过程中出错: {str(e)}", exc_info=True)
        raise

if __name__ == '__main__':
    main()